import numpy as np
import matplotlib.pyplot as plt
a = np.loadtxt('type1/fft_type1.txt',usecols=1)
b = np.loadtxt('type2/fft_type2.txt',usecols=1)

n = 132000
#fourier = np.fft.fft(angles)
timestep = 0.01 # in ps
#freq = np.fft.fftfreq(n, d=timestep)

'''
a_re = np.reshape(a,(n,int(len(a)/n)),order='F') 
a_sum = np.sum(a_re,axis=1)

xf1 = np.linspace(0.0, 1.0//(2.0*timestep), n//2)*4.1357
yf1 = 2.0/n * np.abs(a_sum[:n//2])/(len(a)/n)

b_re = np.reshape(b,(n,int(len(b)/n)),order='F')
b_sum = np.sum(b_re,axis=1)

xf2 = np.linspace(0.0, 1.0//(2.0*timestep), n//2)*4.1357
yf2 = 2.0/n * np.abs(b_sum[:n//2])/(len(b)/n)

np.savetxt('fft_sum1.txt',np.array([xf1,yf1]).T)
np.savetxt('fft_sum2.txt',np.array([xf2,yf2]).T)
'''


a_re = np.reshape(a,(n,int(len(a)/n)),order='F') 

xf1 = np.linspace(0.0, 1.0//(2.0*timestep), n//2)*4.1357
yf1 = 2.0/n * np.abs(a_re[:n//2])/(len(a)/n)
print(yf1.shape)
yf1_sum = np.sum(yf1,axis=1)
print(yf1.shape)

b_re = np.reshape(b,(n,int(len(b)/n)),order='F')

xf2 = np.linspace(0.0, 1.0//(2.0*timestep), n//2)*4.1357
yf2 = 2.0/n * np.abs(b_re[:n//2])/(len(b)/n)
yf2_sum = np.sum(yf2,axis=1)

np.savetxt('fft_ave_sum1.txt',np.array([xf1,yf1_sum]).T)
np.savetxt('fft_ave_sum2.txt',np.array([xf2,yf2_sum]).T)
